package cn.xiaobai.handler.websocket;

import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import cn.xiaobai.domain.Message;
import cn.xiaobai.service.DataPerceptionService;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

/**
 * 通用websocket 服务端
 *
 * @Author yangdaji
 * @Date 2022/11/8 0008
 * @PackageName xiaobai.service.impl.websocket
 * @ClassName GeneralWebsocketServer
 */
@Slf4j
@Component
@ServerEndpoint(value = "/") // 直接写 '/' 可实现通过 ip + 端口连接
public class WsPushDataHandler {

    /**
     * 数据感知服务
     */
    private static DataPerceptionService dataPerceptionService;

    @Autowired
    public void setDataPerceptionService(DataPerceptionService dataPerceptionService) {
        WsPushDataHandler.dataPerceptionService = dataPerceptionService;
    }

    /**
     * 会话地图 key: 频道号, value: 会话对象
     */
    private static final Map<String, Collection<Session>> channel2SessionMap = new ConcurrentHashMap<>();

    /**
     * 连接绑定频道, key: 会话id, value: 频道号
     */
    private static final Map<String, String> session2ChannelMap = new ConcurrentHashMap<>();

    /**
     * 任务执行线程
     */
    private static ThreadPoolTaskScheduler taskScheduler;

    @Autowired
    public void setTaskScheduler(ThreadPoolTaskScheduler taskScheduler) {
        WsPushDataHandler.taskScheduler = taskScheduler;
    }

    /**
     * 频道推送隔离, key: 频道号, value: 推送线程
     */
    private static final Map<String, ThreadPoolTaskScheduler> scheduledTasksPool2Map = new ConcurrentHashMap(16);

    private Session session;

    private ScheduledFuture<?> scheduledFuture;

    /**
     * 间隔周期
     */
    private Integer period;

    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
    }

    @OnMessage
    public void onMessage(String org0, Session session) {
        log.info("接收到订阅数据: {}", org0);
        Message message = JSONUtil.toBean(org0, Message.class);
        // 频道和数据类型
        String channel = message.getChannel();
        period = Objects.isNull(message.getPeriod()) ? 1 : message.getPeriod();

        session2ChannelMap.put(session.getId(), channel);
        final Collection<Session> sessions = channel2SessionMap.get(channel);
        if (!CollectionUtils.isEmpty(sessions)) {
            sessions.add(session);
        } else {
            CopyOnWriteArraySet<Session> sessionsSet = new CopyOnWriteArraySet<>();
            sessionsSet.add(session);
            channel2SessionMap.put(channel, sessionsSet);
            scheduledTasksPool2Map.put(channel, taskScheduler);
        }
        broadCast();
    }

    /**
     * 数据推送
     */
    public void broadCast() {
        if (!CollectionUtils.isEmpty(channel2SessionMap)) {
            // 通过会话绑定的频道进行获取对应会话
            String channel = session2ChannelMap.get(session.getId());
            Collection<Session> sessions = channel2SessionMap.get(channel);

            ThreadPoolTaskScheduler threadPoolTaskScheduler = scheduledTasksPool2Map.get(channel);
            if (Objects.nonNull(threadPoolTaskScheduler)) {
                try {
                    if (Objects.nonNull(scheduledFuture)) {
                        scheduledFuture.cancel(false);
                    }
                    scheduledFuture = threadPoolTaskScheduler.scheduleAtFixedRate(() -> {
                        // 查询最新推送数据
                        Optional<String> jsonMessageOptional = dataPerceptionService
                                .queryPushDataByChannel(channel);
                        if (!jsonMessageOptional.isPresent()) {
                            return;
                        }
                        String jsonMessage = jsonMessageOptional.get();
                        if (!CollectionUtils.isEmpty(channel2SessionMap)) {
                            if (!CollectionUtils.isEmpty(sessions)) {
                                sessions.forEach(sess -> {
                                    try {
                                        //连接正常, 则推送
                                        if (sess.isOpen()) {
                                            sess.getBasicRemote().sendText(jsonMessage);
                                        }
                                    } catch (IOException e) {
                                        log.error("推送失败: {}", e.getMessage());
                                        close();
                                    }
                                });
                            }
                        }
                    }, TimeUnit.SECONDS.toMillis(period));
                } catch (Exception e) {
                    log.error("任务执行失败: {}", e.getMessage());
                    close();
                }
            }
        }
    }

    @OnClose
    public void close() {
        if (Objects.nonNull(session)) {
            try {
                session.close();
            } catch (IOException e) {
                log.error("会话关闭失败: {}", e.getMessage());
            }
            removeSession(session);
        }
    }

    /**
     * 删除会话绑定关系
     *
     * @param session 当前会话
     */
    private void removeSession(Session session) {
        String id = session.getId();
        String channel = session2ChannelMap.get(id);
        // 解除会话对应推送频道
        session2ChannelMap.remove(id);
        // 解除频道对应会话
        if (channel2SessionMap.containsKey(channel)) {
            Collection<Session> sessions = channel2SessionMap.get(channel);
            sessions.stream()
                    .filter(sess -> Objects.equals(sess.getId(), id))
                    .forEach(sess -> {
                        try {
                            sess.close();
                            sessions.remove(sess);
                        } catch (IOException e) {
                            log.error("会话关闭失败: {}", e.getMessage());
                        }
                    });
        }
        if (Objects.nonNull(scheduledFuture)) {
            scheduledFuture.cancel(false);
        }
    }

}
